#include <time.h>
#include <fcntl.h>
#include <netdb.h>
#include <stdio.h>
#include <errno.h>
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <syslog.h>
#include <pthread.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/select.h>
#include <netinet/in.h>
#include "err_quit.h"
#include "http.h"
#include "list.h"
#include "safe_free.h"
#include "include.h"
#include "fix_close_wait.h"
#include "util_debug.h"

extern int h_errno;    /* for get hostbyname #include <netdb.h> */

/* system may failed somtimes, wo need to change signal behavior of SIG_CHLD */
typedef void(*sighandler_t)(int);

/* openssl */
SSL_CTX *ctx_s;
SSL_CTX *ctx_c;
static char *server_cert_file = "./ssl/server.crt";
static char *private_key_file = "./ssl/server.key";
static int l_fd = -1;
static int c_fd = -1;
static char *server_ip = NULL;
static unsigned short server_port = 0;
static char *local_ip = NULL;
static unsigned short local_port = 0;
static int fd_dump_req = -1;
static int fd_dump_rsp = -1;

void sig_handler(int signo);
int ssl_init(void);
void signal_init(void);
int proxy_listen_ssl(void);
void process_ssl(int _fd);
void thread_close_wait(void *arg);
void thread_wait(void *arg);

int sock_set_nonblock(int fd);
int sock_set_reuseaddr(int _fd);
int sock_get_peeraddr(int fd, char *ip, unsigned short *port);
int sock_get_localaddr(int fd, char *ip, unsigned short *port);

int create_proxy_server(char *host, unsigned short l_port, int listen_num);
int create_real_server(const char *host, unsigned short port);
int create_real_server_nonblock(const char *host, unsigned short port, int sec);
void iptables_flush();
void iptables_init();

int main(int argc, char **argv)
{
    if(argc == 2 && (0 == strcmp(argv[1], "flush"))) {
        iptables_flush();
        printf("Flush all HTTPS_DUMP iptables\n");
        return 0;
    }
    if(argc != 5) {
        printf("https_dump server_ip server_port local_ip local_port\n");
        return 0;
    }

    server_ip = argv[1];
    server_port = (unsigned short)atoi(argv[2]);
    local_ip = argv[3];
    local_port = (unsigned short)atoi(argv[4]);
    iptables_flush();
    iptables_init();

    /* 初始化openssl, ctx_s, ctx_c等 */
    if(ssl_init() < 0) {
        fprintf(stderr, "cannot ssl_init()\n");
        //syslog(LOG_INFO, "ssl_init failed");
        return 0;
    }

    /* 建立socket */
    l_fd = create_proxy_server(local_ip, local_port, 100);
    if(l_fd < 0) {
        printf("cannot create proxy server\n");
        syslog(LOG_INFO, "create proxy server failed");
        return 0;
    }

    signal_init();

    /* 子线程:回收子进程 */
    pthread_t th;
    if(0 != pthread_create(&th, NULL, (void *)&thread_wait, NULL)) {
        perror("pthread_create(wait)");
        return 0;
    }

    proxy_listen_ssl();

    return 0;
}

/*
 *  worker:
 *          SIGPIPE: 忽略
 *          SIGHUP: 退出进程
 *  为https设置了一个定时器
 *      用于临时解决openssl造成的close_wait的bug 
 */
void sig_handler(int signo)
{
    switch(signo) {
    case SIGALRM:
        {
#ifdef CLOSE_WAIT
            pthread_t th_close_wait;
            if(0 != pthread_create(&th_close_wait, NULL,
                        (void *)&thread_close_wait, NULL)) {
                perror("pthread_create(close_wait)");
            }
            alarm(TIME_FIX_CLOSE_WAIT);
#endif
            break;
        }
    case SIGUSR1:
        break;
    case SIGUSR2:
        break;
    case SIGPIPE:
        break;
    default:
        exit(1);
    }
}


int ssl_init(void)
{ 
#ifdef FUNC
    printf("==========start ssl_init()==========\n");
#endif
    SSL_load_error_strings();
    //OpenSSL_add_ssl_algorithms();
    SSL_library_init();

    ctx_c = SSL_CTX_new(TLSv1_2_client_method());  //代理客户端
    if(!ctx_c) {
#ifdef DEBUG_SSL
        printf("cannot create ctx_c\n");
#endif
        return -1;
    }
    ctx_s = SSL_CTX_new(TLSv1_2_server_method());  //代理服务器
    if(!ctx_s) {
#ifdef DEBUG_SSL
        printf("cannot create ctx_s\n");
#endif
        return -1;
    }

    //SSL_CTX_set_verify(ctx_s, SSL_VERIFY_NONE, NULL);
    //SSL_CTX_set_verify(ctx_s, SSL_VERIFY_PEER, NULL);
    //SSL_CTX_load_verify_locations(ctx_s, ca_cert_file, NULL);
    if(SSL_CTX_use_certificate_file(ctx_s, server_cert_file, SSL_FILETYPE_PEM) <= 0) {
#ifdef DEBUG_SSL
        printf("cannot load server certificate file\n");
#endif
        return -1;
    }
    if(SSL_CTX_use_PrivateKey_file(ctx_s, private_key_file, SSL_FILETYPE_PEM) <= 0) {
#ifdef DEBUG_SSL
        printf("cannot load server private key file\n");
#endif
        return -1;
    }
    if(!SSL_CTX_check_private_key(ctx_s)) {
#ifdef DEBUG_SSL
        printf("cannot match server_cert_file and private_key_file\n");
#endif
        return -1;
    }
    //SSL_CTX_set_cipher_list(ctx_s, "RC4-MD5");
    //SSL_CTX_set_cipher_list(ctx_s, "AES256-GCM-SHA384");
    SSL_CTX_set_cipher_list(ctx_s, "ALL");
    SSL_CTX_set_mode(ctx_s, SSL_MODE_AUTO_RETRY);
#ifdef FUNC
    printf("==========finish ssl_init()==========\n");
#endif
    return 0;
}

void signal_init(void)
{
    if(signal(SIGINT, sig_handler) == SIG_ERR) {
        err_quit("signal()");
    }
    printf("register SIGINT=%d\n", SIGINT);

    if(signal(SIGSEGV, sig_handler) == SIG_ERR) {
        err_quit("signal()");
    }
    printf("register SIGSEGV=%d\n", SIGSEGV);

    if(signal(SIGPIPE, sig_handler) == SIG_ERR) {
        err_quit("signal()");
    }
    printf("register SIGPIPE=%d\n", SIGPIPE);

    if(signal(SIGUSR1, sig_handler) == SIG_ERR) {
        err_quit("signal()");
    }
    printf("register SIGUSR1=%d\n", SIGUSR1);

    if(signal(SIGUSR2, sig_handler) == SIG_ERR) {
        err_quit("signal()");
    }
    printf("register SIGUSR2=%d\n", SIGUSR2);

    if(signal(SIGUSR2, sig_handler) == SIG_ERR) {
        err_quit("signal()");
    }

    /* SIGALRM用于定时处理回收子进程 */
    if(signal(SIGALRM, sig_handler) == SIG_ERR) {
        err_quit("signal");
    }
    alarm(TIME_FIX_CLOSE_WAIT);

}


int proxy_listen_ssl(void)
{
    printf("\n==========start proxy_listen_ssl(%d)==========\n", getpid());

    struct sockaddr_in client_addr;
    bzero(&client_addr, sizeof(client_addr));
    socklen_t len_client = sizeof(client_addr);
    while(1) {
        c_fd = accept(l_fd, (struct sockaddr *)&client_addr, &len_client);
        if(c_fd < 0) {
            perror("cannot accept correctly, accept()");
            continue;
        }
        pid_t pid = fork();
        switch(pid) {
        case -1:
            close(c_fd);
            perror("fork()");
            break;
        case 0:
            close(l_fd);
            process_ssl(c_fd);
            exit(0);
        default:
            close(c_fd);
            break;
        }
    }
    //隐式回收
    printf("==========finish proxy_listen()==========\n");
    return 0;
}

SSL *ssl_server_handshake(int c_fd)
{
    SSL *ssl_s = SSL_new(ctx_s);
    if(NULL == ssl_s) {
        debug_print("%s: Failed\n", __func__);
        return NULL;
    }
    int ret = SSL_set_fd(ssl_s, c_fd);
    if(ret != 1) {
        print_ssl_error(ssl_s, ret, "ssl_server_handshake: SSL_set_fd()");
        SSL_free(ssl_s);
        return NULL;
    }

    SSL_set_accept_state(ssl_s);

    while(1) {
        int ret = SSL_do_handshake(ssl_s);
        if(1 != ret) {
            int err = print_ssl_error(ssl_s, ret, "ssl_server_handshake: SSL_do_handshake()");
            if(err == SSL_ERROR_WANT_READ || 
               err == SSL_ERROR_WANT_WRITE) {
                continue;
            }
            else {
                SSL_free(ssl_s);
                ssl_s = NULL;
                break;
            }
        }
        break;
    }
    return ssl_s;
}

SSL *ssl_client_handshake(int s_fd)
{
    SSL *ssl_c = SSL_new(ctx_c);
    if(NULL == ssl_c) {
        debug_print("%s: Failed\n", __func__);
        return NULL;
    }

    int ret = SSL_set_fd(ssl_c, s_fd);
    if(ret != 1) {
        print_ssl_error(ssl_c, ret, "ssl_client_handshake: SSL_set_fd()");
        SSL_free(ssl_c);
        return NULL;
    }

    SSL_set_connect_state(ssl_c);
    
    while(1) {
        ret = SSL_do_handshake(ssl_c);
        if(ret != 1) {
            int err = print_ssl_error(ssl_c, ret, "ssl_client_handshake SSL_do_handshake()");
            if(err == SSL_ERROR_WANT_READ ||
               err == SSL_ERROR_WANT_WRITE) {
                continue;
            }
            else {
                SSL_free(ssl_c);
                ssl_c = NULL;
                break;
            }
        }
        break;
    }
    return ssl_c;
}

int ssl_read_forward(int fd_from, SSL *ssl_from, int fd_to, SSL *ssl_to)
{
    int rd = 0;
    int read_blocked = 0;
    int ssl_pending = 0;
    char buff[LEN_SSL_RECORD + 1] = {0};
    while(1) {
        rd = SSL_read(ssl_from, buff, sizeof(buff) - 1);
        ssl_pending = SSL_pending(ssl_from);
        if(ssl_pending > 0) {
            debug_print("%s: SSL_pending %d\n", __func__, ssl_pending);
        }

        int rd_err = print_ssl_error(ssl_from, rd, "ssl_read_forward: SSL_read()");
        switch(rd_err) {
        case SSL_ERROR_NONE:
            {
                int wr = SSL_write(ssl_to, buff, rd);
                int wr_err = print_ssl_error(ssl_to, wr, "ssl_read_forward: SSL_write()");
                //出错处理
            }
            break;
        case SSL_ERROR_WANT_READ:
            /* nothing to read in socket */
            read_blocked = 1;
            break;
        case SSL_ERROR_SYSCALL:
        case SSL_ERROR_ZERO_RETURN:
            return 0;
        default:
            break;
        }

        if(read_blocked) {
            rd = 1;
            break;
        }
        if(ssl_pending == 0) {
            rd = 1;
            break;
        }
    }

    return rd; 
}

void process_ssl(int _fd)
{
    /* https:一问一答模式 */
    printf("%s start\n", __func__);
    int ret;
    int c_fd = _fd;
    int s_fd = -1;
    SSL *ssl_s = NULL;
    SSL *ssl_c = NULL;

    time_t t = time(NULL);
    struct tm tm;
    localtime_r(&t, &tm);
    char request[128] = {0};
    char response[128] = {0};
    char ip[24] = {0};
    unsigned short port = 0;
    get_peer_ip(c_fd, ip, &port);
    snprintf(request, sizeof(request) - 1, "%04d-%02d-%02d_%02d:%02d:%02d-%s:%d.request.http",
            tm.tm_year+1900, tm.tm_mon+1, tm.tm_mday, tm.tm_hour, tm.tm_min, tm.tm_sec, ip, port);
    snprintf(response, sizeof(response) - 1, "%04d-%02d-%02d_%02d:%02d:%02d-%s:%d.response.http",
            tm.tm_year+1900, tm.tm_mon+1, tm.tm_mday, tm.tm_hour, tm.tm_min, tm.tm_sec, ip, port);

    /* ssl */
    sock_set_nonblock(c_fd); 
    ssl_s = ssl_server_handshake(c_fd);

    /* 链接真正服务器 */
    s_fd = create_real_server(server_ip, server_port);
    if(s_fd < 0) {
        goto worker_exit;
    }
    sock_set_nonblock(s_fd);

    ssl_c = ssl_client_handshake(s_fd);
    if(ssl_c == NULL) {
        goto worker_exit;
    }

    fd_dump_req = open(request, O_RDWR | O_CREAT | O_TRUNC, 0666);
    if(fd_dump_req < 0) {
        debug_err("%s: open(%s) %s\n", __func__, request, strerror(errno));
        goto worker_exit;
    }

    fd_dump_rsp = open(response, O_RDWR | O_CREAT | O_TRUNC, 0666);
    if(fd_dump_rsp < 0) {
    
    }

    int max_fd = -1;
    fd_set rset, allset;
    FD_ZERO(&allset);
    FD_SET(c_fd, &allset);
    max_fd = c_fd>max_fd?c_fd:max_fd;
    FD_SET(s_fd, &allset);
    max_fd = s_fd>max_fd?s_fd:max_fd;

    while(1) {
        memcpy(&rset, &allset, sizeof(fd_set));
        struct timeval t = {1, 0};
        int ret = select(max_fd + 1, &rset, NULL, NULL, &t);
        if(ret < 0) {
            if(errno == EINTR) {
                continue;
            }
        }
        else if(0 == ret) {
            continue;
        }
        int ret_c2s = 1, ret_s2c = 1;
        if(FD_ISSET(c_fd, &rset)) {
            ret_c2s = ssl_read_forward(c_fd, ssl_s, s_fd, ssl_c, fd_dump_req);
        }
        if(FD_ISSET(s_fd, &rset)) {
            ret_s2c = ssl_read_forward(s_fd, ssl_c, c_fd, ssl_s, fd_dump_rsp);
        }

        if(ret_c2s <= 0 || ret_s2c <= 0) {
            break;
        }
    }
worker_exit:
    if(c_fd > 0) {
        shutdown(c_fd, SHUT_RDWR);
        close(c_fd);
    }
    if(s_fd > 0) {
        shutdown(s_fd, SHUT_RDWR);
        close(s_fd);
    }

    if(ssl_s) {
        SSL_free(ssl_s);
    }
    if(ssl_c) {
        SSL_free(ssl_c);
    }
}

#ifdef CLOSE_WAIT
/*
 * 临时解决https_proxy的CLOSE_WAIT-bug
 * fix_close_wait()还没仔细检查，存在发生内存泄露的隐患
 */
void thread_close_wait(void *arg)
{
    pthread_detach(pthread_self());
    switch(fork()) {
    case 0:
        close(l_fd);
        close(c_fd);
        fix_close_wait();
        exit(0);
        break;
    case -1:
    default:
        break;
    }
    pthread_exit(NULL);
}
#endif

/*
 * 线程不会被子进程继承,哪个线程调用fork，哪个线程被继承
 */
void thread_wait(void *arg)
{
    pthread_detach(pthread_self());
    while(1) {
        /* 阻塞调用 */
        if(wait(NULL) < 0) {
            if(ECHILD == errno) {
                /* 还没有子进程的时候wait立即返-1，会导致CPU占用了50% */
                sleep(1);
            }
        }
    }
    pthread_exit(NULL);
}

/*
 * return:
 * failed : -1
 * ok     : old_flags
 */
int sock_set_nonblock(int fd)
{
    int flags = fcntl(fd, F_GETFL, 0);
    if(flags < 0) {
        debug_err("%s: fcntl(F_GETFL), %s\n", __func__, strerror(errno));
        return -1;
    }
    if(fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) {
        debug_err("%s: fcntl(F_SETFL), %s\n", __func__, strerror(errno));
        return -1;
    }
    return flags;
}

/*
 * return:
 *   0  : ok
 *   -1 : failed
 */
int sock_set_reuseaddr(int _fd)
{
    int opt = 1;
    if(setsockopt(_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        debug_err("%s: setsockopt(), %s\n", __func__, strerror(errno));
        return -1;
    }
    return 0;
}

/*
 * return:
 *  failed: -1
 *  ok    : 0
 * get noting, if ip == NULL && port == NULL
 * only get ip: if port == NULL
 * only get port, if ip == NULL
 * get both ip and port
 */
int sock_get_peeraddr(int fd, char *ip, unsigned short *port)
{
    struct sockaddr_in sock;
    socklen_t len = sizeof(sock);
    if(getpeername(fd, (struct sockaddr *)&sock, &len) < 0) {
        debug_err("%s: getpeername(), %s\n", __func__, strerror(errno));
        return -1;
    }
    if(ip) {
        strcpy(ip, inet_ntoa(sock.sin_addr));
    }
    if(port) {
        *port = ntohs(sock.sin_port);
    }
    return 0;
}


/*
 * return:
 *  failed: -1
 *  ok    : 0
 * get noting, if ip == NULL && port == NULL
 * only get ip: if port == NULL
 * only get port, if ip == NULL
 * get both ip and port
 */
int sock_get_localaddr(int fd, char *ip, unsigned short *port)
{
    struct sockaddr_in sock;
    socklen_t len = sizeof(sock);
    if(getsockname(fd, (struct sockaddr *)&sock, &len) < 0) {
        debug_err("%s: getsockname(), %s\n", __func__, strerror(errno));
        return -1;
    }
    if(ip) {
        strcpy(ip, inet_ntoa(sock.sin_addr));
    }
    if(port) {
        *port = ntohs(sock.sin_port);
    }
    return 0;
}

int create_proxy_server(char *host, unsigned short l_port, int listen_num)
{
#ifdef FUNC
    printf("==========start create_proxy_server()==========\n");
#endif
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if(fd < 0)
        err_quit("socket");
    int opt = 1;
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    struct sockaddr_in local_addr;
    memset(&local_addr, 0, sizeof(local_addr));

    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(l_port);
    if(NULL == host) {
        local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    } 
    inet_pton(AF_INET, host, &local_addr.sin_addr.s_addr);
    if(bind(fd, (struct sockaddr *) &local_addr, sizeof(local_addr)) < 0)
        err_quit("bind");
    if(listen(fd, listen_num) < 0)
        err_quit("listen");
#ifdef FUNC
    printf("==========finish create_proxyy_server()==========\n");
#endif
    return fd;
}

int create_real_server(const char *host, unsigned short port)
{
    /* 建立和服务器的连接, 使用select超时连接 */
#ifdef FUNC
    printf("==========start create_real_server()==========\n");
#endif
#ifdef DEBUG
    printf("create_real_server host=%s, port=%d\n", host, port);
#endif
    int s_fd = socket(AF_INET, SOCK_STREAM, 0);
    if(s_fd < 0) {
        perror("socket()");
        return -1;
    }
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    server_addr.sin_addr.s_addr = inet_addr(host);

    if(connect(s_fd, (struct sockaddr *) &server_addr, sizeof(server_addr)) < 0) {
        perror("connect");
        close(s_fd);
        return -1;
    }
#ifdef DEBUG
    printf("connected to %s:%d\n", host, port);
#endif
#ifdef FUNC
    printf("==========finish create_real_server()==========\n");
#endif
    return s_fd;
}

int create_real_server_nonblock(const char *host, unsigned short port, int sec)
{
    /* 建立和服务器的连接 */
#ifdef FUNC
    printf("==========start create_real_server_nonblock()==========\n");
#endif
    int s_fd = socket(AF_INET, SOCK_STREAM, 0);

    if(s_fd < 0) {
        perror("socket()");
        return -1;
    }
    /* 设置非阻塞 */
    int flags = fcntl(s_fd, F_GETFL, 0);
    if(flags < 0)
    {
        perror("fcntl f_get");
        goto end;
    }
    if(fcntl(s_fd, F_SETFL, flags | O_NONBLOCK) < 0)
    {
        perror("fcntl f_set");
        goto end;
    }

    struct sockaddr_in server_addr;
    struct hostent *server;
    if((server = gethostbyname(host)) == NULL)
    {
        printf("\033[31m");
        printf("gethostbyname [%s] error, h_error=%d, %s\n", host, h_errno, hstrerror(h_errno));
        printf("\033[0m");
        goto end;
    }
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    /* inet_pton(AF_INET, host, &(server_addr.sin_addr.s_addr)); */
    memcpy(&(server_addr.sin_addr.s_addr), server->h_addr, server->h_length);
#ifdef DEBUG
    char ip[16] = {0};
    printf("%s <--> %s port=%d\n", host, inet_ntop(AF_INET, server->h_addr, ip, sizeof(ip)), port);
#endif
    if(connect(s_fd, (struct sockaddr *) &server_addr, sizeof(server_addr)) < 0)
    {
        if(errno != EINPROGRESS)
        {
#ifdef DEBUG
            syslog(LOG_INFO, "cannot connect to %s:%d", inet_ntop(AF_INET, server->h_addr, ip, sizeof(ip)), port);
            printf("connect err\n");
#endif
            goto end;
        }
    }
    fd_set rset, wset;
    FD_ZERO(&rset);
    FD_ZERO(&wset);
    FD_SET(s_fd, &rset);
    FD_SET(s_fd, &wset);
    struct timeval tout;
    tout.tv_sec = sec > 0 ? sec : 0;
    tout.tv_usec = 0;
    int ret = select(s_fd + 1, &rset, &wset, NULL, tout.tv_sec > 0 ? &tout : NULL);
    if(ret > 0)
    {
        if(FD_ISSET(s_fd, &rset) || FD_ISSET(s_fd, &wset))
        {
            int error = 0;
            unsigned int len = sizeof(error);
            if(getsockopt(s_fd, SOL_SOCKET, SO_ERROR, &error, &len) < 0)
            {
                perror("getsockopt");
                goto end;
            }
            else
            {
                /* 改回非阻塞 */
                if(fcntl(s_fd, F_SETFL, flags) < 0)
                    goto end;
                return s_fd;
            }
        }
    }
    else if(ret == 0)
    {
#ifdef DEBUG
        printf("select timeout!\n");
#endif
        goto end;
    }
    else
    {
        perror("select");
        goto end;
    }

end:
    close(s_fd);
#ifdef FUNC
    printf("==========finish create_real_server_nonblock()==========\n");
#endif
    return -1;
}

int SYSTEM(const char *format, ...)
{
	char buf[4096]={0};
	va_list arg;
	int ret = 0;
	memset(buf, 0, sizeof(buf));
	va_start(arg, format);
	vsnprintf(buf,4096, format, arg);
	va_end(arg);

    printf("%s\n", buf);
	system(buf);

	return ret;
}

void iptables_flush()
{
    SYSTEM("iptables -t nat -D PREROUTING -j HTTPS_DUMP");
    SYSTEM("iptables -t nat -F HTTPS_DUMP");
    SYSTEM("iptables -t nat -X HTTPS_DUMP");
}


void iptables_init()
{
    SYSTEM("iptables -t nat -N HTTPS_DUMP");
    SYSTEM("iptables -t nat -I PREROUTING -j HTTPS_DUMP");
    SYSTEM("iptables -t nat -A HTTPS_DUMP -p tcp -d %s --dport %d -j DNAT --to-destination %s:%d",
            server_ip, server_port, local_ip, local_port);
}

